// created by lupo
// created at 2017/07/21

#include <vector>
#include <thread>
#include <iostream>
#include <unistd.h>
#include <cuda_runtime.h>  
#include <device_launch_parameters.h> 


void worker(int thread_id, int gpu_id)
{
    int gpu_id_tmp;
    cudaError_t ret;
    ret = cudaGetDevice(&gpu_id_tmp);
    if (ret != cudaSuccess) {
        std::cout << "cudaGetDevice failed, ret: " << ret << std::endl;
        return;
    } else {
        std::cout << "before set, thread_id: " << thread_id 
                  << ", gpu_id:" << gpu_id_tmp << std::endl;
    }

    ret = cudaSetDevice(gpu_id);
    if (ret != cudaSuccess) {
        std::cout << "cudaSetDevice failed, ret: " << ret << std::endl;
        return;
    }

    ret = cudaGetDevice(&gpu_id_tmp);
    if (ret != cudaSuccess) {
        std::cout << "after set, cudaGetDevice failed, ret: " << ret << std::endl;
        return;
    } else {
        std::cout << "after set, thread_id: " << thread_id 
                  << ", gpu_id:" << gpu_id_tmp << std::endl;
    }
    sleep(5);
}


int main(int argc, char* argv[])
{
    int gpu_cnt = 2;
    const int thread_cnt = 10;
    std::vector<std::thread> threads;
    if (argc > 1) {
        gpu_cnt = std::stoi(argv[1]);
    }

    for (int i=0; i<thread_cnt; ++i) {
        threads.push_back(std::thread(worker, i+1, i%gpu_cnt));
        sleep(2);
    }

    for (int i=0; i<thread_cnt; ++i) {
        threads[i].join();
    }

    return 0;
}
